from __future__ import print_function
import argparse
import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
import time
from models import *
from PIL import Image
from utils.config import *
from utils import config

parser = argparse.ArgumentParser(description='PSMNet')
parser.add_argument('--datatype', default='kitti2015', help='datapath')
parser.add_argument('--loadmodel', default='./save/kitti2015/kitti.tar', help='loading model')
parser.add_argument('--model', default='stackhourglass', help='select model')
parser.add_argument('--kfold', type=int, default=0, help='Kfold num', choices=range(6)) # 5折交叉验证，若为0则不使用交叉验证
parser.add_argument('--maxdisp', default=192, help='maxium disparity')
parser.add_argument('--no-cuda', action='store_true', default=False, help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

if args.datatype == 'kitti2015':
    from dataloader import KITTI2015LoaderTest as DA
    if system == 'Windows': datapath = 'G:/KITTI/KITTI2015/data_scene_flow/training/'
    elif system == 'Ubuntu': datapath = '/media/ubuntu/e/zhouyiqing/KITTI/KITTI2015/data_scene_flow/training/'
elif args.datatype == 'kitti2012':
    from dataloader import KITTI2012LoaderTest as DA
    if system == 'Windows': datapath = 'G:/KITTI/KITTI2012/data_stereo_flow/training/'
    elif system == 'Ubuntu': datapath = '/media/ubuntu/e/zhouyiqing/KITTI/KITTI2012/data_stereo_flow/training/'
elif args.datatype == 'usvinland':
    if args.kfold != 0: config.kfold = args.kfold
    from dataloader import USVInlandLoaderTest as DA
    if system == 'Windows': datapath = 'G:/USVInland/Stereo Matching/Low_Res_640_320/'
    elif system == 'Ubuntu': datapath = '/media/ubuntu/e/zhouyiqing/USVInland/Stereo Matching/Low_Res_640_320/'
elif args.datatype == 'usvinland_seg':
    if args.kfold != 0: config.kfold = args.kfold
    from dataloader import USVInlandLoaderSegTest as DA
    if system == 'Windows': datapath = 'G:/USVInland/Stereo Matching/Segmentation/'
    elif system == 'Ubuntu': datapath = '/media/ubuntu/e/zhouyiqing/USVInland/Stereo Matching/Segmentation/'

test_left_img, test_right_img = DA.dataloader(datapath)

if args.model == 'stackhourglass':
    model = stackhourglass(args.maxdisp)
elif args.model == 'basic':
    model = basic(args.maxdisp)
else:
    print('no model')

model = nn.DataParallel(model, device_ids=[0])
model.cuda()

if args.loadmodel is not None:
    state_dict = torch.load(args.loadmodel)
    model.load_state_dict(state_dict['state_dict'])

print('Number of model parameters: {}'.format(sum([p.data.nelement() for p in model.parameters()])))

def test(imgL,imgR):
    model.eval()

    if args.cuda:
        imgL = imgL.cuda()
        imgR = imgR.cuda()

    with torch.no_grad():
        output = model(imgL,imgR)
    output = torch.squeeze(output).data.cpu().numpy()
    return output

def main():
    normal_mean_var = {
        'mean': [0.485, 0.456, 0.406],
        'std': [0.229, 0.224, 0.225]}
    infer_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(**normal_mean_var)])

    # # 预热，让模型先跑50轮
    # x, y = torch.randn(1, 3, 640, 640).cuda(), torch.randn(1, 3, 640, 640).cuda()
    # for i in range(50):
    #     output = model(x, y)

    for inx in range(len(test_left_img)):
        imgL_o = Image.open(test_left_img[inx]).convert('RGB')
        imgR_o = Image.open(test_right_img[inx]).convert('RGB')

        imgL = infer_transform(imgL_o)
        imgR = infer_transform(imgR_o)

        # pad to width and hight to 16 times
        if imgL.shape[1] % 16 != 0:
            times = imgL.shape[1]//16
            top_pad = (times+1)*16 -imgL.shape[1]
        else:
            top_pad = 0

        if imgL.shape[2] % 16 != 0:
            times = imgL.shape[2]//16
            right_pad = (times+1)*16-imgL.shape[2]
        else:
            right_pad = 0

        imgL = F.pad(imgL,(0,right_pad, top_pad,0)).unsqueeze(0)
        imgR = F.pad(imgR,(0,right_pad, top_pad,0)).unsqueeze(0)

        torch.cuda.synchronize()
        start_time = time.time()
        pred_disp = test(imgL,imgR)
        torch.cuda.synchronize()
        print('inx = %d, time = %.6f' %(inx, time.time() - start_time))

        if top_pad !=0 or right_pad != 0:
            img = pred_disp[top_pad:,:-right_pad]
        else:
            img = pred_disp

        img = img.astype('uint16')
        # img = Image.fromarray(img)
        # img.save(os.path.join('./out', args.datatype, test_left_img[inx].split('/')[-1]))

        if 'usvinland' in args.datatype: img = img * 255.0 / 50.0
        img_np = img.astype(np.uint8).squeeze()
        disp_img = cv2.cvtColor(img_np, cv2.COLOR_BAYER_BG2GRAY) # 灰度图像
        disp_pseudo = cv2.applyColorMap(cv2.convertScaleAbs(disp_img, alpha=5), cv2.COLORMAP_JET) # 映射到伪彩色，近红远蓝

        # 保存预测视差和伪彩色图像
        out_path = os.path.join('./out', args.datatype)
        if not os.path.exists(os.path.join(out_path, 'pseudo')):
            os.makedirs(os.path.join(out_path, 'pseudo'))
        cv2.imwrite(os.path.join(out_path, test_left_img[inx].split('/')[-1]), disp_img)
        cv2.imwrite(os.path.join(out_path, 'pseudo', test_left_img[inx].split('/')[-1]), disp_pseudo)

if __name__ == '__main__':
    main()
